import sys, re
sys.path.append('/datas/projects/python_backends/src')
from utils.db import client
import logging, json, requests
import time
_client = client()


def predict(data):
 
    _params = json.loads(data['params'])

    body = dict()
    body['version'] = _client.get_version(_params['model'])
    body['input'] = _params


    logging.info(f'模型实际入参: {body}')

    ## 发出请求
    url = 'https://api.replicate.com/v1/predictions'
    headers = {"Authorization":"Token r8_cxpdRLfLyWo2CbJf7SEh5XPPEftcKk62nGZeo", "Content-Type":"application/json"}

    res = requests.post(url, headers=headers, data=json.dumps(body))
    resdata = json.loads(res.content)
    logging.info(f'模型返回:{resdata}')
    print(f'模型返回:{resdata}')

    ## 更新数据库，关联taskid和progressId
    task_id = data['task_id']
    progress_id = resdata.get("id")
    logging.info(f'connect task_id: {task_id} ==> progress_id {progress_id}')
    _client.update_sd_progressid(task_id, progress_id)

    c = 0
    while(True):
        ## 获取进度
        progress = get_progress(task_id)
        _client.update_sd_job_progress(task_id, progress)

        if progress == -1.0:
            return "failed"
        if progress == 1.0:
            return 

        time.sleep(1)
        c += 1
          
        if c > 60:
            return "failed"
 
def get_progress(task_id):
    def get_res(progress_id):
        url = f'https://api.replicate.com/v1/predictions/{progress_id}'
        headers = {"Authorization":"Token r8_cxpdRLfLyWo2CbJf7SEh5XPPEftcKk62nGZeo", "Content-Type":"application/json"}

        ## 获取task对应的progressid
        res = requests.get(url, headers=headers)
        res = json.loads(res.content)
        logging.info(f'进程查看:{res}')
        print(f'进程查看:{res}')

        return res


    progress_id = _client.get_progress(task_id)
    if not progress_id:
        return -1.0
        logging.error('获取progress_id失败')
    logging.info(f'get progress_id {progress_id} from {task_id}')
    print(f'get progress_id {progress_id} from {task_id}')

    res = get_res(progress_id)
    ## 开始处理后播报
    nt = 1
    while(res.get('status') == 'starting'):
        time.sleep(nt)
        nt += 1
        res = get_res(progress_id)

    progress = re.findall('\d+%', res.get('logs', ''))
    current_progress = [0.01*float(i.replace('%',  '')) for i in progress] + [0.0]
    fseed = re.findall('seed:\s*(-?\d+)', res.get('logs', ''))
    seed = fseed[0] if fseed else -1

    if res.get('output') and res.get('status') == 'succeeded' :
        ## 生成文件后，获取并上传oss
        tmp = list()
        for item in res.get('output'):
            tmp.append(_client.uploadfile(item))
        res['output'] = tmp

    ## 队列中，未开始
    if res.get('status') == 'starting' and not res.get('logs'):
        time.sleep(1)

    ## 失败直接返回
    if res.get('status') == 'failed':
        return -1
    
    if max(current_progress) == 1.0 and res.get('status') == 'succeeded':
        _client.update_image(task_id, res.get('output')[0])
        return 1.0
    else:
        return max(current_progress)





    
if __name__ == '__main__':
    params = {
        'task_id':'3gpsufbA6siE5Hb', 
        'params': json.dumps({"model": "bytedance/sdxl-lightning-4step", "image": "", "width": 1024, "height": 1024, "sampler": "Euler a", "seed": -1, "steps": 30, "cfg_scale": 7, "face_fix": False, "hd_fix": False, "hd_redraw_rate": 0.5, "hd_scale": 2, "hd_scale_alg": "Latent", "hd_steps": 15, "prompt": "1girl", "negative_prompt": "nsfw, paintings,low quality,easynegative,ng_deepnegative ,lowres,bad anatomy,bad hands,bad feet", "session_id": "x3k8viasrrd0eyrrn99shm0rey82vu1dpynjeq9gef"})
    }
    predict(params)